"""
数据标注模块
使用大模型对数据进行标注
"""
import pandas as pd
import json
import requests
import logging
from typing import Dict, List, Any, Optional
import time
import asyncio
import aiohttp
from datetime import datetime

logger = logging.getLogger(__name__)

class DataLabeler:
    """数据标注器"""
    
    def __init__(self, model_config: Dict[str, Any]):
        """
        初始化数据标注器
        
        Args:
            model_config: 模型配置，包含API endpoint, key等信息
        """
        self.model_config = model_config
        self.api_endpoint = model_config.get('api_endpoint')
        self.api_key = model_config.get('api_key')
        self.model_name = model_config.get('model_name', 'gpt-3.5-turbo')
        self.max_retries = model_config.get('max_retries', 3)
        self.retry_delay = model_config.get('retry_delay', 1)
        
    def create_labeling_prompt(self, data_row: pd.Series, labeling_rules: Dict[str, Any]) -> str:
        """
        创建标注提示词
        
        Args:
            data_row: 数据行
            labeling_rules: 标注规则
            
        Returns:
            标注提示词
        """
        # 获取数据内容
        data_content = data_row.to_dict()
        
        # 构建提示词
        prompt = f"""
请对以下数据进行标注分析：

数据内容：
{json.dumps(data_content, ensure_ascii=False, indent=2)}

标注规则：
{json.dumps(labeling_rules, ensure_ascii=False, indent=2)}

请根据标注规则对数据进行分析，并以JSON格式返回标注结果。
返回格式示例：
{{
    "labels": {{
        "category": "分类结果",
        "sentiment": "情感分析结果",
        "quality_score": 评分,
        "tags": ["标签1", "标签2"]
    }},
    "confidence": 0.95,
    "reasoning": "标注理由"
}}
"""
        return prompt
    
    def call_llm_api(self, prompt: str) -> Dict[str, Any]:
        """
        调用大模型API进行标注
        
        Args:
            prompt: 提示词
            
        Returns:
            标注结果
        """
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        
        data = {
            'model': self.model_name,
            'messages': [
                {
                    'role': 'user',
                    'content': prompt
                }
            ],
            'temperature': 0.1,
            'max_tokens': 1000
        }
        
        for attempt in range(self.max_retries):
            try:
                response = requests.post(
                    self.api_endpoint,
                    headers=headers,
                    json=data,
                    timeout=30
                )
                
                if response.status_code == 200:
                    result = response.json()
                    content = result['choices'][0]['message']['content']
                    
                    # 尝试解析JSON
                    try:
                        return json.loads(content)
                    except json.JSONDecodeError:
                        # 如果不是JSON格式，返回原始内容
                        return {
                            'labels': {'raw_response': content},
                            'confidence': 0.8,
                            'reasoning': 'Raw LLM response'
                        }
                else:
                    logger.error(f"API调用失败: {response.status_code} - {response.text}")
                    
            except Exception as e:
                logger.error(f"API调用异常 (尝试 {attempt + 1}/{self.max_retries}): {str(e)}")
                
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay * (2 ** attempt))  # 指数退避
                else:
                    # 最后一次尝试失败，返回默认结果
                    return {
                        'labels': {'error': 'API调用失败'},
                        'confidence': 0.0,
                        'reasoning': f'API调用失败: {str(e)}'
                    }
        
        return {
            'labels': {'error': 'API调用失败'},
            'confidence': 0.0,
            'reasoning': 'Maximum retries exceeded'
        }
    
    def label_single_row(self, data_row: pd.Series, labeling_rules: Dict[str, Any]) -> Dict[str, Any]:
        """
        标注单行数据
        
        Args:
            data_row: 数据行
            labeling_rules: 标注规则
            
        Returns:
            标注结果
        """
        prompt = self.create_labeling_prompt(data_row, labeling_rules)
        result = self.call_llm_api(prompt)
        
        # 添加元数据
        result['timestamp'] = datetime.now().isoformat()
        result['model'] = self.model_name
        
        return result
    
    def label_dataframe(self, df: pd.DataFrame, labeling_rules: Dict[str, Any], 
                       batch_size: int = 10) -> pd.DataFrame:
        """
        标注整个数据框
        
        Args:
            df: 数据框
            labeling_rules: 标注规则
            batch_size: 批处理大小
            
        Returns:
            包含标注结果的数据框
        """
        results = []
        total_rows = len(df)
        
        logger.info(f"开始标注 {total_rows} 行数据")
        
        for i in range(0, total_rows, batch_size):
            batch_df = df.iloc[i:i+batch_size]
            
            for idx, row in batch_df.iterrows():
                try:
                    result = self.label_single_row(row, labeling_rules)
                    results.append(result)
                    
                    logger.info(f"完成标注 {len(results)}/{total_rows} 行")
                    
                except Exception as e:
                    logger.error(f"标注第 {idx} 行失败: {str(e)}")
                    results.append({
                        'labels': {'error': f'标注失败: {str(e)}'},
                        'confidence': 0.0,
                        'reasoning': f'标注过程中发生错误: {str(e)}',
                        'timestamp': datetime.now().isoformat(),
                        'model': self.model_name
                    })
            
            # 批次间休息
            if i + batch_size < total_rows:
                time.sleep(1)
        
        # 将结果添加到原数据框
        df_labeled = df.copy()
        df_labeled['labeling_result'] = results
        
        # 提取标注结果到单独的列
        df_labeled['labels'] = [r['labels'] for r in results]
        df_labeled['confidence'] = [r['confidence'] for r in results]
        df_labeled['reasoning'] = [r['reasoning'] for r in results]
        df_labeled['labeling_timestamp'] = [r['timestamp'] for r in results]
        
        logger.info(f"标注完成，共处理 {len(df_labeled)} 行数据")
        
        return df_labeled
    
    def save_labeled_data(self, df_labeled: pd.DataFrame, output_path: str):
        """
        保存标注结果
        
        Args:
            df_labeled: 标注后的数据框
            output_path: 输出路径
        """
        try:
            # 保存为CSV
            if output_path.endswith('.csv'):
                df_labeled.to_csv(output_path, index=False, encoding='utf-8')
            # 保存为JSON
            elif output_path.endswith('.json'):
                df_labeled.to_json(output_path, orient='records', 
                                 force_ascii=False, indent=2)
            # 保存为Excel
            elif output_path.endswith('.xlsx'):
                df_labeled.to_excel(output_path, index=False)
            else:
                # 默认保存为CSV
                df_labeled.to_csv(output_path + '.csv', index=False, encoding='utf-8')
            
            logger.info(f"标注结果已保存到: {output_path}")
            
        except Exception as e:
            logger.error(f"保存标注结果失败: {str(e)}")
            raise

# 使用示例
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 模型配置
    model_config = {
        'api_endpoint': 'https://api.openai.com/v1/chat/completions',
        'api_key': 'your-api-key',
        'model_name': 'gpt-3.5-turbo',
        'max_retries': 3,
        'retry_delay': 1
    }
    
    # 标注规则
    labeling_rules = {
        'categories': ['正面', '负面', '中性'],
        'quality_criteria': {
            'relevance': '内容相关性',
            'accuracy': '准确性',
            'completeness': '完整性'
        },
        'tag_options': ['重要', '紧急', '常规', '低优先级']
    }
    
    # 示例数据
    sample_data = pd.DataFrame({
        'text': ['这是一条正面评论', '这是一条负面评论', '这是一条中性评论'],
        'source': ['用户A', '用户B', '用户C'],
        'timestamp': ['2024-01-01', '2024-01-02', '2024-01-03']
    })
    
    # 创建标注器
    labeler = DataLabeler(model_config)
    
    # 进行标注
    labeled_data = labeler.label_dataframe(sample_data, labeling_rules)
    
    # 保存结果
    labeler.save_labeled_data(labeled_data, 'labeled_results.csv')
    
    print("标注完成!") 